# 在线版面分析
import cv2
import numpy as np
from common.get_hw_pre import getresult
import json
import os
from tqdm import tqdm
# 数据集文件夹
data_dir = "E:\internship\dataset_test/format_new/9.1"
# ===================================================================
# 在线还是离线, 在线需要内网，离线是使用保存的json预测结果
on_line = False
# 标签格式是否是云平台,不然就是是线下的
label_style = False
# 如果是离线需要输入保存的json文件
pre_json_dir = os.path.join(data_dir,"hw_bo_result.txt")

# 日志文件夹
img_savedir = os.path.join(data_dir,"format_det_result")
os.makedirs(img_savedir,exist_ok=True)
class_name_dict = {"0":"picture","1":"paragraph","2":"table","3":"cell"}
if not on_line:
    fpre_json_dir = open(pre_json_dir,"r",encoding="utf-8")
    pre_jsons = fpre_json_dir.readlines()


filenames = []
anns = []

# 不同的标签格式加载方式不一样
if not label_style:
    for line in open(os.path.join(data_dir, "Label.txt"), "r", encoding='utf-8'):
        strs = line.split("\t")
        filename = strs[0]

        # filename = filename.split(".")[0] + ".jpg"
        # path = os.path.join(data_dir, filename)
        path = os.path.join(data_dir, "../" + filename)
        path = path.replace("\\","/")
        filenames.append(path)

        anns.append(json.loads(strs[1]))
else:
    for line in open(os.path.join("E:/internship/dataset_test/exportjson_16.39/bo_file", "2022-10-24.txt"), "r", encoding='utf-8'):
        strs = line.split("：")
        if not len(strs) == 3:
            print("error: please check label format")
        filename = strs[1].split(" ")[0]
        # filename = filename.split(".")[0] + ".jpg"
        path = os.path.join(data_dir, "../" + filename)
        path = path.replace("\\", "/")
        filenames.append(path)
        anns.append(json.loads(strs[2]))

sample_metrics = []
class_labels = []
for i, _ in enumerate(tqdm(filenames)):
    label = anns[i]

    class_label = []
    class_pre = []
    class_pre_boxs = []
    class_label_boxs = []
    #label
    if not label_style:
        for j in range(len(label)):
            value = label[j]["transcription"]
            # if value == 'picture':
            if 'picture' in value:
                class_label.append(0)
                xyxy = label[j]["points"]
                class_label_boxs.append(xyxy)
            elif 'para' in value or 'Para' in value:
                class_label.append(1)
                xyxy = label[j]["points"]
                class_label_boxs.append(xyxy)

            elif 'table' in value:
                class_label.append(2)
                xyxy = label[j]["points"]
                class_label_boxs.append(xyxy)

            elif value.split("-")[0] == 'cell':
                class_label.append(3)
                xyxy = label[j]["points"]
                class_label_boxs.append(xyxy)
    else:
        for current_label in label[0]["Labels"]:
            value = current_label['Value']
            xyxy = current_label['Points']
            point1 = [round(xyxy[0]["X"]), round(xyxy[0]["Y"])]
            point2 = [round(xyxy[1]["X"]), round(xyxy[1]["Y"])]
            point3 = [round(xyxy[2]["X"]), round(xyxy[2]["Y"])]
            point4 = [round(xyxy[3]["X"]), round(xyxy[3]["Y"])]
            if 'picture' in value:
                class_label.append(0)
                class_label_boxs.append([point1,point2,point3,point4])
            elif 'para' in value or 'Para' in value:
                class_label.append(1)
                class_label_boxs.append([point1,point2,point3,point4])
            elif 'table' in value:
                class_label.append(2)
                class_label_boxs.append([point1,point2,point3,point4])

            elif value.split("-")[0] == 'cell':
                class_label.append(3)
                class_label_boxs.append([point1,point2,point3,point4])
    # pre:
    if on_line:
        pre = getresult(filenames[i])
    else:
        pre = json.loads(pre_jsons[i].split("\t")[1])
    if 'Picture' in pre.keys():
        for class_info in pre["Picture"]:
            coords = class_info["coords"]
            class_pre_boxs.append(coords)
            class_pre.append(0)
    if 'Paraph' in pre.keys():
        for class_info in pre["Paraph"]:
            coords = class_info["coords"]
            class_pre_boxs.append(coords)
            # class_pre_boxs.append([coords[0], coords[1], coords[4], coords[5]])
            class_pre.append(1)
    if 'Table' in pre.keys():
        for class_info in pre["Table"]:
            coords = class_info["coords"]
            class_pre_boxs.append(coords)
            # class_pre_boxs.append([coords[0], coords[1], coords[4], coords[5]])
            class_pre.append(2)
            for cell in class_info["cell"]:
                class_pre.append(3)
                coords = cell["coords"]
                class_pre_boxs.append(coords)
                # class_pre_boxs.append([coords[0], coords[1], coords[4], coords[5]])
    from common.detect import get_per_img_statistics
    current_metrics = get_per_img_statistics(class_pre_boxs, class_label_boxs, 0.5, class_pre, class_label)
    sample_metrics += current_metrics
    class_labels += class_label

    # 分析并画图
    unique_classes = np.unique(class_label).tolist() + np.unique(class_pre).tolist()
    unique_classes = set(unique_classes)
    imgcopy = -1
    for class_name_index, class_name in enumerate( unique_classes):
        # 在预测结果中 class_name 类别的编号
        current_class_pre_index = np.where(np.array(class_pre) == class_name)[0]
        # 在标签中 class_name 类别的编号
        current_class_gt_index = np.where(np.array(class_label) == class_name)[0]
        # 如果个数不一样，则肯定有漏检或者误检的，则需要画图
        len_not_equal = not (len(current_class_pre_index) == len(current_class_gt_index))
        tp_exist0 = False
        if len(current_class_pre_index) * len(current_class_gt_index) == 0 :
            tp_exist0 = True
        else:
            # 如果个数一样，那么看一看计算出的TP结果，是否有0存在，如果存在那么需要画图，否则不画图
            if 0 in current_metrics[class_name_index][0]:
                tp_exist0 = True
        if tp_exist0 or len_not_equal:
            img = cv2.imread(filenames[i])
            if img is None:
                print("!!!!!!error cv2.imread() plase check img dir")
            imgcopy = img.copy()
            # 绘制预测框
            pre_boxes = np.array(class_pre_boxs)[current_class_pre_index]
            for prebox_index, pre_box in enumerate(pre_boxes):
                # 默认用绿色
                color = [0, 255, 0]
                # 如果该类别没有真实框，那么用红色
                if len(current_class_gt_index) == 0:
                    color = [0, 0, 255]
                else:
                    # 如果当前框的TP=0 用红色
                    tp_zeros = np.where(current_metrics[class_name_index][0] == 0)[0]
                    if len(tp_zeros) > 0:
                        if prebox_index in tp_zeros:
                            color = [0, 0, 255]
                # pre_box 是8个数值
                for point_index in range(int(len(pre_box)/2)):
                    offset = point_index * 2
                    start = [pre_box[0 + offset],pre_box[1 + offset]]
                    # 从最后一个点画到第一个点，数组内存溢出
                    if 2 + offset > 7 :
                        offset = -2
                    end = [pre_box[2 + offset],pre_box[3 + offset]]

                    cv2.line(imgcopy, start, end, color, 2)

            labe_boxes = np.array(class_label_boxs)[current_class_gt_index]
            for labe_box in labe_boxes:
                # labe_box 是4个二维数组
                for point_index in range(int(len(labe_box))):
                    offset = point_index
                    start = labe_box[0 + offset]
                    if 1 + offset > 3 :
                        offset = -1
                    end = labe_box[1 + offset]
                    cv2.line(imgcopy, start, end, [255,0,0], 2)
            save_dir = os.path.join(img_savedir, class_name_dict[str(class_name)])
            os.makedirs(save_dir,exist_ok=True)
            cv2.imwrite(os.path.join(save_dir, filenames[i].split("/")[-1]), imgcopy)

from common.detect import ap_per_class
true_positives, pred_scores, pred_labels = [np.concatenate(x, 0) for x in list(zip(*sample_metrics))]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives, pred_scores, pred_labels, class_labels)

mAP = AP.mean()
print("precision ",precision)
print("recall ",recall)
print("AP ",AP)
print("f1 ",f1)
print("mAP ",mAP)